// Model selection

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <math.h>

#include <set>
#include <map>

#include "itf_debug.h"

#include "prefix_tree.h"
#include "alg_ms.h"
#include "mex.h"

#define LOGLK_SIZE 3

using namespace std;

namespace model
{

/* The likelihood ratio calculation */
/*-----------------------------------------------------------------------*/
static double loglikelihood_exp(double *clkarray, unsigned int clk_size, double lambda) {
    unsigned int i;
    double sum = 0;
    
    for(i=0; i<clk_size; i++)
    {
        sum += log(lambda*exp(-lambda*clkarray[i]));
    }
    
    return sum;
}

static double loglikelihood_wei(double *clkarray, unsigned int clk_size, double gamma, double k) {
    unsigned int i;
    double sum = 0;
    
    for(i=0; i<clk_size; i++)
    {
        sum += log((k/gamma)*pow(clkarray[i]/gamma,k-1)*exp(-pow(clkarray[i]/gamma, k)));
    }
    
    return sum;
}

static double loglikelihood_logn(double *clkarray, unsigned int clk_size, double mu, double sigma) {
    mxArray *input[2], *out[1];
    double *clk, *al;
    
    input[0] =  mxCreateDoubleMatrix(1, 2, mxCOMPLEX);
    input[1] =  mxCreateDoubleMatrix(1, clk_size, mxCOMPLEX);
    
    al = (double *) mxGetData(input[0]);
    clk = (double *) mxGetData(input[1]);
    
    //Copy clkarray to clk
    memcpy(clk, clkarray, clk_size*sizeof(double));
    
    al[0] = mu;
    al[1] = sigma;
    
    /* Call the Matlab funtion 'expfit(data,alpha)' */
    mexCallMATLAB(1, out, 2, input, "lognlike");
    
    return -mxGetScalar(out[0]);
    
}
/*-----------------------------------------------------------------------*/

/* Start matlab Funtions                                                 */
/* --------------------------------------------------------------------- */
/* Estimate exponential parameters */
static void inf_exp(double *clkarray, unsigned int clk_size, double alpha, double *lambda) {
    mxArray *input[2], *out[1];
    double *clk, *al;
    
    input[0] =  mxCreateDoubleMatrix(1, clk_size, mxCOMPLEX);
    input[1] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    
    clk = (double *) mxGetData(input[0]);
    al = (double *) mxGetData(input[1]);
    
    //Copy clkarray to clk
    memcpy(clk, clkarray, clk_size*sizeof(double));
    
    *al = alpha;
    
    /* Call the Matlab funtion 'expfit(data,alpha)' */
    mexCallMATLAB(1, out, 2, input, "expfit");
    
    *lambda = mxGetScalar(out[0]);
}

/* Estimate weibull parameters */
static void inf_wei(double *clkarray, unsigned int clk_size, double alpha,
        double *gamma, double *k) {
    
    mxArray *input[2], *out[1];
    double *clk, *al, *param;
    unsigned int i;
    
    input[0] =  mxCreateDoubleMatrix(1, clk_size, mxCOMPLEX);
    input[1] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    
    clk = (double *) mxGetData(input[0]);
    al = (double *) mxGetData(input[1]);
    
    //Copy clkarray to clk
    //memcpy(clk, clkarray, clk_size*sizeof(double));
    for(i=0; i<clk_size; i++) {
        clk[i]=clkarray[i];
        
        mexPrintf("%f ", clkarray[i]);
    }
    mexPrintf("\n");
    
    *al = alpha;
    
    /* Call the Matlab funtion 'expfit(data,alpha)' */
    mexCallMATLAB(1, out, 2, input, "wblfit");
    
    param = (double *) mxGetData(out[0]);
    
    *gamma = param[0];
    *k = param[1];
}

/* Estimate Lognormal distribution parameters */
static void inf_logn(double *clkarray, unsigned int clk_size, double alpha,
        double *mu, double *sigma) {
    
    mxArray *input[2], *out[1];
    double *clk, *al, *param;
    unsigned int i;
    
    input[0] =  mxCreateDoubleMatrix(1, clk_size, mxCOMPLEX);
    input[1] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    
    clk = (double *) mxGetData(input[0]);
    al = (double *) mxGetData(input[1]);
    
    //Copy clkarray to clk
    //memcpy(clk, clkarray, clk_size*sizeof(double));
    for(i=0; i<clk_size; i++) {
        clk[i]=clkarray[i];
        
        mexPrintf("%f ", clkarray[i]);
    }
    mexPrintf("\n");
    
    *al = alpha;
    
    /* Call the Matlab funtion 'expfit(data,alpha)' */
    mexCallMATLAB(1, out, 2, input, "lognfit");
    
    param = (double *) mxGetData(out[0]);
    
    *mu = param[0];
    *sigma = param[1];
}

/* [h,p] = chi2gof_exp(x, 0.1, 0.5) */
/*static void chi2_exp(clock_l *clkList, double alpha, double lambda,
        double *h, double *p) {
    
    int i;
    mxArray *input[3], *out[2];
    double *clk, *al, *lam;
    clock_n *cl;
    
    input[0] =  mxCreateDoubleMatrix(1, clkList->n, mxCOMPLEX);
    input[1] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    input[2] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    
    clk = (double *) mxGetData(input[0]);
    al = (double *) mxGetData(input[1]);
    lam = (double *) mxGetData(input[2]);
    
    // Convert clock List to Matlab array
    cl = clkList->firstNode;
    for(i=0; i<clkList->n; i++) {
        clk[i] = (double) cl->data;
        cl = cl->next;
    }
    
    *al = alpha;
    *lam = lambda;
    
    // Call the Matlab funtion 'expfit(data,alpha)'
    mexCallMATLAB(2, out, 3, input, "chi2gof_exp");
    
    *h = (double) mxGetScalar(out[0]);
    *p = (double) mxGetScalar(out[1]);
}*/

/*static void chi2_wei(clock_l *clkList, double alpha, double gamma, double k,
        double *h, double *p) {
    
    int i;
    mxArray *input[4], *out[2];
    double *clk, *al, *gam, *mk;
    clock_n *cl;
    
    input[0] =  mxCreateDoubleMatrix(1, clkList->n, mxCOMPLEX);
    input[1] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    input[2] =  mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    input[3] = mxCreateDoubleMatrix(1, 1, mxCOMPLEX);
    
    clk = (double *) mxGetData(input[0]);
    al = (double *) mxGetData(input[1]);
    gam = (double *) mxGetData(input[2]);
    mk = (double *) mxGetData(input[3]);
    
    // Convert clock List to Matlab array
    cl = clkList->firstNode;
    for(i=0; i<clkList->n; i++) {
        clk[i] = (double) cl->data;
        cl = cl->next;
    }
    
    *al = alpha;
    *gam = gamma;
    *mk = k;
    
    // Call the Matlab funtion 'expfit(data,alpha)'
    mexCallMATLAB(2, out, 4, input, "chi2gof_wei");
    
    *h = (double) mxGetScalar(out[0]);
    *p = (double) mxGetScalar(out[1]);
}*/

/* --------------------------------------------------------------------- */
/* End matlab Funtions                                                   */

static unsigned int max_el(double *loglk) {
    unsigned int m, max;
    
    max = 0;
    for(m=1; m<LOGLK_SIZE; m++) {
        if (loglk[m] > loglk[max])
        	max = m;
    }
    
    return max;
}

void clock_debug(double *clocks, unsigned int size, unsigned int event) {
    unsigned int i;
    
    DEBUG_PRINT((FDEBUG, "Event %d: ", event));
    for(i=0; i<size; i++)
        DEBUG_PRINT((FDEBUG, "%f ", clocks[i]));
    DEBUG_PRINT((FDEBUG, "\n"));
}

void initialize_distribution_structure(psa *trie, unsigned int size) {
    trie->parameters = (struct distribution *) malloc(sizeof(struct distribution)*size);
}

void estimate(psa *trie, double *clocks, unsigned int clk_size, unsigned int event, double alpha) {
    
    double exp_lambda, wei_gamma, wei_k, logn_mu, logn_sigma, loglk[LOGLK_SIZE];
    
    inf_exp(clocks, clk_size, alpha, &exp_lambda);
    inf_wei(clocks, clk_size, alpha, &wei_gamma, &wei_k);
    inf_logn(clocks, clk_size, alpha, &logn_mu, &logn_sigma);

    loglk[0] = loglikelihood_exp(clocks, clk_size, exp_lambda);
    loglk[1] = loglikelihood_wei(clocks, clk_size, wei_gamma, wei_k);
    loglk[2] = loglikelihood_logn(clocks, clk_size, logn_mu, logn_sigma);

    DEBUG_PRINT((FDEBUG, "Estimated parameters: Exp(%f):%f, Wei(%f,%f):%f\n", exp_lambda, loglk[0], wei_gamma, wei_k, loglk[1]));

    // find max of loglk array
    switch(max_el(loglk)) {
        case 0:
            trie->parameters[event-1].dist = 0;
            trie->parameters[event-1].dist_parameters[0] = exp_lambda;
            break;
        case 1:
            trie->parameters[event-1].dist = 1;
            trie->parameters[event-1].dist_parameters[0] = wei_gamma;
            trie->parameters[event-1].dist_parameters[1] = wei_k;
            break;
        case 2:
            trie->parameters[event-1].dist = 2;
            trie->parameters[event-1].dist_parameters[0] = logn_mu;
            trie->parameters[event-1].dist_parameters[1] = logn_sigma;
    }
}

/* Select the best distributions for clock data */
void selection(psa *trie) {
    unsigned int i;
    ntrieNode *node_succ;
    
    double alpha = 0.001;
    
    unsigned int pivot, size;
    double *clocks;
    
    set<unsigned int> set_events;
    set<unsigned int>::iterator it_se;
    
    multimap<unsigned int,ntrieNode *> map_event2samples;
    multimap<unsigned int,ntrieNode *>::iterator it_me2s;
    
    for(i=0; i<trie->n; i++) {
        if(trie->nodeArray[i]->removed == 0) {
            // for each active event on trie->nodeArray[i]
            NEDTRIE_FOREACH(node_succ, ntrie_s, trie->nodeArray[i]->child) {
                map_event2samples.insert(pair<unsigned int,ntrieNode *>((unsigned int)node_succ->key, node_succ));
                set_events.insert((unsigned int)node_succ->key);
            }
        }
    }
    
    initialize_distribution_structure(trie, set_events.size());
    
    for(it_se=set_events.begin(); it_se != set_events.end(); it_se++) {
        size=0;
        for(it_me2s = map_event2samples.find(*it_se); it_me2s != map_event2samples.upper_bound(*it_se); it_me2s++)
            size += (*it_me2s).second->clks_size;
        
        DEBUG_PRINT((FDEBUG, "Size: %d, E: %d \n", size, *it_se));
        
        clocks = (double *) malloc(sizeof(double)*size);
        
        pivot=0;
        for(it_me2s = map_event2samples.find(*it_se); it_me2s != map_event2samples.upper_bound(*it_se); it_me2s++) {
            
            //clock_debug((*it_me2s).second->clks, (*it_me2s).second->clks_size, *it_se);
            DEBUG_PRINT((FDEBUG, "--- [%d:%d] size: %d | Event: %d \n", pivot, pivot+((*it_me2s).second->clks_size), size, *it_se));
            
            memcpy(&(clocks[pivot]), (*it_me2s).second->clks, (*it_me2s).second->clks_size*sizeof(double));
            pivot += (*it_me2s).second->clks_size;
        }
        
        // Get parameters from each clock sample set
        clock_debug(clocks, size, *it_se);
        
        estimate(trie, clocks, size, *it_se, alpha);
    }
    
    return;
    
    /*for(i=0; i<trie->n; i++) {
        if(trie->nodeArray[i]->removed == 0) {
            
            // for each sucessor nodes
            NEDTRIE_FOREACH(node_succ, ntrie_s, trie->nodeArray[i]->child) {
                
            }
        }
    }*/
    
}

}

